import os
import sys
import shutil
import subprocess
import argparse

def modify_config(original_cfg, new_cfg, factor_prefix, value, is_enum_mode):
    with open(original_cfg) as fin, open(new_cfg, 'w') as fout:
        for line in fin:
            if line.startswith(factor_prefix + ",") or (is_enum_mode and line.startswith("partition,")):
                parts = line.strip().split(",")
                if is_enum_mode:
                    if parts[2] == factor_prefix.split(",")[2]:
                        parts[1] = value
                        fout.write(",".join(parts) + "\n")
                    else:
                        fout.write(line)
                else:
                    if value == 0:
                        continue  # skip line
                    parts[-1] = str(value)
                    fout.write(",".join(parts) + "\n")
            else:
                fout.write(line)

def is_enum_factor(factor_str):
    try:
        int(factor_str.strip().split(",")[-1])
        return False
    except ValueError:
        return True

def main():
    parser = argparse.ArgumentParser(description="Run Aladdin across factor values and plot results")
    parser.add_argument("--function", required=True, help="Kernel name, e.g. stencil")
    parser.add_argument("--factor", required=True, help="Factor prefix, e.g. partition,cyclic,orig")
    parser.add_argument("--values", required=True, help="Comma-separated list of factor values")
    parser.add_argument("--skip", action="store_true", help="Only generate configs and run Aladdin, skip plotting")
    parser.add_argument("--cached", action="store_true", help="Use cached results, only generate plot")
    args = parser.parse_args()

    aladdin_home = os.environ.get("ALADDIN_HOME")
    if not aladdin_home:
        print("Error: ALADDIN_HOME is not set", file=sys.stderr)
        sys.exit(1)

    example_dir = os.path.join(aladdin_home, "SHOC", args.function, "example")
    trace_path = os.path.join(aladdin_home, "SHOC", args.function, "dynamic_trace.gz")
    orig_cfg = os.path.join(example_dir, "config_example")
    if not os.path.isdir(example_dir) or not os.path.isfile(trace_path) or not os.path.isfile(orig_cfg):
        print("Error: example directory or files missing", file=sys.stderr)
        sys.exit(1)

    tmp_dir = os.path.join(example_dir, "tmp")
    if not args.cached:
        if os.path.exists(tmp_dir):
            shutil.rmtree(tmp_dir)
        os.makedirs(tmp_dir)

    is_enum = is_enum_factor(args.factor)
    values = args.values.split(",") if is_enum else [int(v) for v in args.values.split(",")]

    if not args.cached:
        for v in values:
            run_dir = os.path.join(tmp_dir, str(v))
            os.makedirs(run_dir)
            cfg_path = os.path.join(run_dir, "config_example")
            modify_config(orig_cfg, cfg_path, args.factor, v, is_enum)
            os.symlink(trace_path, os.path.join(run_dir, "dynamic_trace.gz"))
            cmd = [os.path.join(aladdin_home, "common", "aladdin"),
                   args.function,
                   "dynamic_trace.gz",
                   "config_example"]
            print("Running:", " ".join(cmd), "in", run_dir)
            with open(os.path.join(run_dir, "stencil_summary"), "w") as out:
                subprocess.run(cmd, cwd=run_dir, stdout=out, check=True)

    if not args.skip:
        plot_dirs = [os.path.join(tmp_dir, str(v)) for v in values]
        plot_cmd = ["python", "plot_metrics.py"] + plot_dirs + ["--factor", args.factor]
        print("Plotting with:", " ".join(plot_cmd))
        subprocess.run(plot_cmd, check=True)

if __name__ == "__main__":
    main()
